import numpy as np

def compute_plane_normal(depth_image, mask):
    """
    计算掩膜区域对应的深度图平面的法向量。

    Args:
        depth_image (numpy.ndarray): 深度图，二维numpy数组。
        mask (numpy.ndarray): 掩膜，二值numpy数组，大小与depth_image相同。

    Returns:
        numpy.ndarray: 法向量，长度为3的numpy数组，表示平面法向量。
    """
    # 提取掩膜区域的深度值
    points = np.column_stack(np.where(mask == 1))  # 获取掩膜区域的坐标 (u, v)
    z_values = depth_image[mask == 1]  # 提取掩膜区域的深度值

    if points.shape[0] < 3:
        return None  # 少于3个点无法计算平面

    # 将 (u, v, z) 点构成 3D 点云
    points_3d = np.hstack((points, z_values.reshape(-1, 1)))  # (u, v, z)

    # 最小二乘法拟合平面 Ax + By + Cz + D = 0
    A = points_3d[:, 0]  # u (x)
    B = points_3d[:, 1]  # v (y)
    C = points_3d[:, 2]  # z (z)

    # 构建矩阵 A * [A, B, C] = D，其中 D 是常数项
    A_matrix = np.vstack([A, B, np.ones_like(A)]).T
    C_matrix = C.reshape(-1, 1)

    # 使用最小二乘法拟合平面
    plane_params, _, _, _ = np.linalg.lstsq(A_matrix, C_matrix, rcond=None)

    # 平面法向量是拟合结果的系数 [A, B, C]
    normal = np.array([plane_params[0].item(), plane_params[1].item(), -1])

    # 归一化法向量
    norm = np.linalg.norm(normal)
    normal_normalized = normal / norm

    return normal_normalized

if __name__ == "__main__":

    # 示例：假设有一个深度图和掩膜
    depth_image = np.ones((480, 640))  # 示例深度图，值范围是0到1
    mask = np.zeros_like(depth_image)  # 掩膜，假设这里全是非掩膜区域
    mask[100:200, 100:200] = 1  # 设置掩膜区域

    # 计算平面法向量
    normal_vector = compute_plane_normal(depth_image, mask)

    # 输出法向量
    if normal_vector is not None:
        print("平面法向量：", normal_vector)
    else:
        print("掩膜区域点数不足，无法计算法向量")
